#pragma once

#include <vector>
#include <stdint.h>
#include "tensor.h"

class SoftMax
{
public:
    SoftMax() = default;
    virtual ~SoftMax() = default;

    void forward(const std::vector<double>& input, std::vector<double> &output);
    std::vector<std::vector<double>> backward();
    std::vector<std::vector<double>> backward(const std::vector<double> &output);

    void forward(const std::vector<std::vector<double>>& samples, std::vector<std::vector<double>> &outputs);
    std::vector<std::vector<double>> backward(const std::vector<std::vector<double>> &output_derivative);
private:
    std::vector<double> input_;
    std::vector<double> output_{};

    std::vector<std::vector<double>> output_batch_{};
};
